import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

import Patches3D
import prob
from NN import cdinn2



# def gen_plot_data():
#     plot_sample_x = torch.linspace(prob.domain_min[0], prob.domain_max[0], prob.PLOT_LEN_B[0])
#     plot_sample_y = torch.linspace(prob.domain_min[1], prob.domain_max[1], prob.PLOT_LEN_B[1])
#     grid_sample = [plot_sample_x, plot_sample_y]
#     # grid_sample = torch.tensor(grid_sample)
#     mesh = torch.meshgrid(grid_sample)
#     flatten = [torch.flatten(mesh[i]) for i in range(len(mesh))] # flatten the list of meshes
#     plot_data = torch.stack(flatten, 1) # stack the list of flattened meshes
#     return plot_data
#
#
# def plot_boundary(model): # barrier boundary: contour plotting
#     barrier_plot_nn_input = gen_plot_data()
#     # apply the nn model but do not require gradient
#     with torch.no_grad():
#         output = model(barrier_plot_nn_input)
#         barrier_plot_nn_output = output.reshape(prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[0]) # y_size * x_size
#     plot_Z = barrier_plot_nn_output.numpy()
#
#     plot_sample_x = np.linspace(prob.domain_min[0], prob.domain_max[0], prob.PLOT_LEN_B[0])
#     plot_sample_y = np.linspace(prob.domain_min[1], prob.domain_max[1], prob.PLOT_LEN_B[1])
#     plot_X, plot_Y = np.meshgrid(plot_sample_x, plot_sample_y)
#     #plt.contourf(plot_X, plot_Y, plot_Z, [0], color='k')
#     barrier_contour = plt.contour(plot_X.T, plot_Y.T, plot_Z, [0], \
#         linewidths = 1, colors='r')
#     # plt.clabel(barrier_contour, fontsize=20, colors='b')
#     return barrier_contour


def simulation(model):
    fig = plt.figure(
        figsize=(3, 3),
        dpi=200
    )
    ax = fig.add_subplot(1, 1, 1, projection='3d')

    # ==========================init and unsafe==============================
    if prob.init_shape == 1:  # rectangle
        center1 = (prob.init_max[0] + prob.init_min[0]) / 2.0
        center2 = (prob.init_max[1] + prob.init_min[1]) / 2.0
        center3 = (prob.init_max[2] + prob.init_min[2]) / 2.0
        diff1 = center1 - prob.init_min[0]
        diff2 = center2 - prob.init_min[1]
        diff3 = center3 - prob.init_min[2]
        Patches3D.Rectangle_3D(ax, (center1, center2, center3), (diff1, diff2, diff3), facecolors="green", alpha=0.8)
    if prob.init_shape == 2:  # circle
        center1 = (prob.init_max[0] + prob.init_min[0]) / 2.0
        center2 = (prob.init_max[1] + prob.init_min[1]) / 2.0
        center3 = (prob.init_max[2] + prob.init_min[2]) / 2.0
        r = center1 - prob.init_min[0]
        Patches3D.Circle_3D(ax, (center1, center2, center3), r, facecolor='green')

    if prob.unsafe_shape == 1:
        center1 = (prob.unsafe_max[0] + prob.unsafe_min[0]) / 2.0
        center2 = (prob.unsafe_max[1] + prob.unsafe_min[1]) / 2.0
        center3 = (prob.unsafe_max[2] + prob.unsafe_min[2]) / 2.0
        diff1 = center1 - prob.unsafe_min[0]
        diff2 = center2 - prob.unsafe_min[1]
        diff3 = center3 - prob.unsafe_min[2]
        Patches3D.Rectangle_3D(ax, (center1, center2, center3), (diff1, diff2, diff3), facecolors="red", alpha=0.8)
    if prob.unsafe_shape == 2:
        center1 = (prob.unsafe_max[0] + prob.unsafe_min[0]) / 2.0
        center2 = (prob.unsafe_max[1] + prob.unsafe_min[1]) / 2.0
        center3 = (prob.unsafe_max[2] + prob.unsafe_min[2]) / 2.0
        r = center1 - prob.unsafe_min[0]
        Patches3D.Circle_3D(ax, (center1, center2, center3), r, facecolor='red')

    # ==================================0水平集==================================
    boundary = Patches3D.plot_boundary(ax, model)
    # boundary = Patches3D.plot_boundary_withoutModel(ax)

    # ==================================轨迹==================================
    num = 5
    step = 1000
    Patches3D.plot_tran(ax, num, step)

    ax.set_box_aspect([1, 1, 1])
    ax.set_xlim([prob.domain_min[0], prob.domain_max[0]])
    ax.set_ylim([prob.domain_min[1], prob.domain_max[1]])
    ax.set_zlim([prob.domain_min[2], prob.domain_max[2]])
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    plt.show()


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    # model.load_state_dict(torch.load('./model/cdinn_c9_1_20_0.1_0.05_1_train.pt'), strict=True)
    # model.load_state_dict(torch.load('./model/cdinn_c9_1_10_0.1_0.05_1_train.pt'), strict=True)  # biggerC9
    model.load_state_dict(torch.load('./model/cdinn_c9_1_5_0.1_0.1_1_train.pt'), strict=True)  # global bigger c9
    # model.load_state_dict(torch.load('./model/cdinn_c8_1_20_0.1_0.1_1_train.pt'), strict=True)
    # model.load_state_dict(torch.load('./model/cdinn_eg5_1_10_0.1_0.05_0.5_train.pt'), strict=True)
    # model.load_state_dict(torch.load('./model/cdinn_zhrR1_1_20_0.1_0.1_1_train.pt'), strict=True)
    simulation(model)
    # simulation()
    pass